import random

import numpy as np

k = 0.1
α = 0
β = 1


class Env:
    def __init__(self):
        self.observation_space_shape = (1, 40)
        self.obs = np.zeros((2, 4, 5))
        self.actions = self.initActionSpace(8, 2)
        self.action_space_num = 256
        self.mec_capality = 350
        self.mec_computeablity = 25
        self.mec_uelist = []
        self.mec_uenum = [0, 0]
        self.transmission_rate = 10
        self.user_c_local = 8
        self.local_consume = []
        self.rewards = 0
        self.time = 0

    # 初始化环境，构思如何用数组描述该环境
    def reset(self):
        """
            obs具有以下结构：
                [
                    mec1: [ue1: [], ue2: [], ue3: [], ue4: []],
                    mec2: [ue1: [], ue2: [], ue3: [], ue4: []]
                ]
                整个环境包含两个mec：
                    mec = obs[i]
                每个mec最多包含4个用户：
                    ue = mec[i]
                每个用户具有以下结构：
                    ue = [M, F, T, R, S]
                    其中M表示用户任务计算量，F表示用户任务所需计算资源，T表示用户最长等待时延,S表示任务是否正在计算
                    若M = 0 代表此用户不存在，当Si != 0时表示任务正在计算，此时针对uei的动作不生效，该用户占用的资源仍保留，只计算其他用户的消耗
                    The size of the task is randomly assigned at the range of [100, 150] MB
                    and fi = εMi is the required computation resource, where ε is randomly set from [0.4, 0.5] Gcycles/MB.
                    Tmax ~ (10s, 20s)
                e.g.:
                    [
                        [
                            [100, 40, 80, 0],
                            [0, 0, 0, 0],
                            [130, 50, 90, 0],
                            [0, 0, 0, 0],
                        ],
                        [
                            [100, 40, 80, 0],
                            [0, 0, 0, 0],
                            [130, 50, 90, 0],
                            [0, 0, 0, 0],
                        ]
                    ]

                MEC性能：
                    计算容量C


        """
        self.mec_uelist = []
        self.mec_uenum = [0, 0]
        self.rewards = 0
        self.time = 0
        for i in range(0, 2):
            for j in range(0, 4):
                if np.random.randint(0, 2) == 1:
                    self.obs[i][j] = self.inituser()
                    self.mec_uelist.append(1)
                    self.mec_uenum[i] += 1
                else:
                    self.obs[i][j] = np.array([0, 0, 0, 0, 0])
                    self.mec_uelist.append(0)

        # print(self.obs.reshape((1, 40)))
        return self.obs.reshape((1, 40))

    # 生成用户数组
    def inituser(self):
        M = round(random.uniform(100, 150), 2)
        ε = round(random.uniform(0.4, 0.5), 1)
        F = ε * M
        T = round(random.uniform(10, 20), 2)
        user = [M, F, T, 0, 0]
        return np.array(user)

    # 环境根据传入的动作返回下一个状态、奖励和是否结束
    def step(self, action):
        # chose_action包含8个元素，对应每个用户的决策
        chose_action = self.actions[action]

        # 判断用户数量是否超出mec服务器容量
        check = self.check_overflow(chose_action)
        if check or self.time == 64:
            next_obs = self.obs.copy()
            self.rewards = 0
            dones = True
        else:
            # 如果没有溢出 更新当前状态
            self.time += 1
            next_obs, dones = self.update_obs(chose_action)

        return next_obs, self.rewards, dones

    def check_overflow(self, action):
        taskcontent = 0
        for i in range(0, 2):
            for j in range(0, 4):
                # 如果动作为1（卸载）并且该用户存在，判断任务总量是否溢出
                if (action[i * 4 + j] == 1) & (self.mec_uelist[i * 4 + j] == 1):
                    taskcontent += self.obs[i][j][0]

                    # 如果mec分配给任务j的计算资源在需求时间内计算完任务，则返回溢出失败
                    if (self.obs[i][j][2] / (self.mec_computeablity / self.mec_uenum[i])) > self.obs[i][j][2]:
                        return True

            if taskcontent > self.mec_capality:
                return True

        return False

    # 生成动作空间
    def initActionSpace(self, uenum, chose):
        result = []
        for i in range(0, 256):
            b = []
            num = [0, 0, 0, 0, 0, 0, 0, 0]
            n = i
            while True:
                s = n // 2  # 商
                y = n % 2  # 余数
                b = b + [y]  # 每一个余数存储到b中

                if s == 0:
                    break  # 余数为0时结束循环

                n = s
            for j in range(0, len(b)):
                num[j] += b[j]

            result.append(num)

        return result

    # 得到本地计算能耗
    def get_local(self):
        consume = []

        # 对每个用户任务进行更新
        for i in range(0, 2):
            for j in range(0, 4):
                # 如果该用户存在，查看其状态
                if self.mec_uelist[i * 4 + j] == 1:
                    if self.obs[i][j][3] == 0:

                        Tl = self.obs[i][j][1] / self.user_c_local
                        El = k * self.obs[i][j][1] * self.user_c_local * self.user_c_local

                        Cli = α * Tl + β * El

                        consume.append(Cli)
                        continue

                    # 如果为其他值，说明仍在计算，不更新状态
                    else:
                        consume.append(0)
                        continue
                else:
                    consume.append(0)
                    continue

        self.local_consume = consume

    def update_obs(self, action):
        dones = False

        consume = []
        self.get_local()

        # 对每个用户任务进行更新
        for i in range(0, 2):
            for j in range(0, 4):
                # 如果该用户是卸载状态而动作选择了本地计算，直接返回
                if (action[i * 4 + j] != self.obs[i][j][4]) and (self.obs[i][j][3] >= 1):
                    consume.append(0)
                    self.obs[i][j][3] -= 1
                    # rewards = -1000
                    continue

                # 如果该用户不存在，生成一个用户添加到状态中
                if self.mec_uelist[i * 4 + j] == 0:
                    consume.append(0)
                    if np.random.randint(0, 2) == 1:
                        self.obs[i][j] = self.inituser()
                        self.mec_uelist[i * 4 + j] = 1
                        self.mec_uenum[i] += 1
                    else:
                        self.obs[i][j] = np.array([0, 0, 0, 0, 0])
                    continue

                # 如果该用户存在，查看其状态
                if self.mec_uelist[i * 4 + j] == 1:
                    # 如果选择了卸载
                    if action[i * 4 + j] == 1:
                        # 如果其任务剩余量 == 1，说明计算完成，重新生成一个用户添加到状态中
                        if self.obs[i][j][3] == 1:
                            consume.append(0)
                            if np.random.randint(0, 2) == 1:
                                self.obs[i][j] = self.inituser()
                            else:
                                self.obs[i][j] = np.array([0, 0, 0, 0, 0])
                                self.mec_uelist[i * 4 + j] = 0
                                self.mec_uenum[i] -= 1
                            continue
                        # 如果任务剩余量 == 0 说明没有开始计算，不更新其状态，计算
                        elif self.obs[i][j][3] == 0:
                            self.obs[i][j][3] = int(self.obs[i][j][1] / (self.mec_computeablity / self.mec_uenum[i]))
                            self.obs[i][j][4] = 1

                            Tmec = self.obs[i][j][1] / (self.mec_computeablity / self.mec_uenum[i]) + \
                                   self.obs[i][j][0] / self.transmission_rate
                            Emec = 20 * self.obs[i][j][0] / self.transmission_rate
                            Cli = α * Tmec + β * Emec

                            consume.append(Cli)

                            continue

                        # 如果为其他值，说明仍在计算，不更新状态
                        else:
                            self.obs[i][j][3] -= 1
                            consume.append(0)
                            continue
                    # 如果本地计算
                    else:
                        # 如果其任务剩余量 == 1，说明计算完成，重新生成一个用户添加到状态中
                        if self.obs[i][j][3] == 1:
                            consume.append(0)
                            if np.random.randint(0, 2) == 1:
                                self.obs[i][j] = self.inituser()
                            else:
                                self.obs[i][j] = np.array([0, 0, 0, 0, 0])
                                self.mec_uelist[i * 4 + j] = 0
                                self.mec_uenum[i] -= 1
                            continue
                        # 如果任务剩余量 == 0 说明没有开始计算，不更新其状态，计算
                        elif self.obs[i][j][3] == 0:
                            self.obs[i][j][3] = int(self.obs[i][j][1] / self.user_c_local)
                            self.obs[i][j][4] = 0
                            Tl = self.obs[i][j][1] / self.user_c_local
                            El = k * self.obs[i][j][1] * self.user_c_local * self.user_c_local

                            Cli = α * Tl + β * El

                            consume.append(Cli)

                            """
                                计算消耗：
                                
                            """
                            continue

                        # 如果为其他值，说明仍在计算，不更新状态
                        else:
                            self.obs[i][j][3] -= 1
                            consume.append(0)
                            continue

        self.get_reward(self.local_consume, consume)
        # print("local:  ******************")
        # print(self.local_consume)
        # print("consume:  ******************")
        # print(consume)

        next_obs = self.obs.copy().reshape((1, 40))
        return next_obs, dones

    def get_reward(self, local_consume, consume):
        reward = 0
        for i in range(0, 8):
            if local_consume[i] != 0:
                # reward += (local_consume[i] - consume[i]) / local_consume[i]
                reward += 0.01 * (local_consume[i] - consume[i])
                # print(reward)
        self.rewards += reward


if __name__ == '__main__':

    env = Env()
    env.reset()

    print(env.step(200))
